﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Text;
using SharpML.Api.Extensions;
using OpenCL.Net;
using SharpML.Api.Implementation;
using SharpML.Api.Interfaces;
using SharpML.Api.Types;
using SharpML.Api.Enumerations;

namespace SharpML.Engines.DeepBelief.Implementation
{
    public class RbmProgram : IMLComputation
    {
        #region - Private Member Variables -

            // Computation Objects
            private InterLayerWeights _weights;
            private List<Layer> _inputs;

            // Native Memory
            private float[] _epsilon = new float[1];
            private float[] _detectorError;
            private float[] _reconstructionError;
            private float[] _model;
            private float[] _reconstructed;
            private float[] _reconstructedModel;
            private float[] _twoPhaseBuffer;
            private int _testCount;

            // Computation Memory
            private List<IMem> _testCaseMems;
            private IMem _weightsMem;
            private IMem _detectorErrorMem;
            private IMem _epsilonMem;
            private IMem _modelMem;
            private IMem _reconstructedMem;
            private IMem _reconstructionErrorMem;
            private IMem _reconstructedModelMem;
            private IMem _errorSumsMem;
            private IMem _twoPhaseMem;

            // Functions
            private Kernel _activateLowerToUpper;
            private Kernel _activateLowerToUpperBinary;
            private Kernel _activateUpperToLower;
            private Kernel _accumulateErrors;
            private Kernel _sumErrorsDetector;
            private Kernel _sumErrorsRecon;
            private Kernel _clearDetectorErrors;
            private Kernel _updateWeights;

	    #endregion

        #region - Construction -

            public RbmProgram(InterLayerWeights weights, List<Layer> inputs)
            {
                _weights = weights;
                _inputs = inputs;
                _testCaseMems = new List<IMem>();

                _model = new float[weights.UpperLayerSize];
                _reconstructed = new float[weights.LowerLayerSize];
                _reconstructionError = new float[_weights.Weights.Length];
                _reconstructedModel = new float[weights.LowerLayerSize];
                _twoPhaseBuffer = new float[weights.LowerLayerSize > weights.UpperLayerSize ? weights.LowerLayerSize : weights.UpperLayerSize];
                _testCount = inputs.Count;
                _detectorError = new float[weights.UpperLayerSize * weights.LowerLayerSize];
            }
 
	    #endregion

        #region - Private Functionality  -

            private string Source
            {
                get
                {
                    StringBuilder sb = new StringBuilder();

                    foreach( var path in Directory.GetFiles( Path.GetDirectoryName( Assembly.GetExecutingAssembly().Location.ResolveRelativePath( "./OpenCL" ) ), "*.rbm.c" ) )
                        sb.Append( File.ReadAllText( path ) );

                    // TODO: REmove need for counts
                    return sb.ToString()
                            .Replace( "<UpperCount>", "1000" )
                            .Replace( "<LowerCount>", "784" );
                }
            }
 
	    #endregion

        #region - Public Functionality  -

            public void Compile( ComputationEngine engine )
            {
                string[] kernelNames = new string [] { 
                    "ActivateLowerToUpper_Binary", 
                    "ActivateLowerToUpperBinary_Binary", 
                    "ActivateUpperToLower", 
                    "AccumulateErrors", 
                    "SumErrorsDetector", 
                    "SumErrorsRecon", 
                    "ClearDetectorErrors", 
                    "UpdateWeights" 
                };

                var kernels = engine.Compile( new [] { ComputationEngine.PsuedoRandomNumbers, Source }, kernelNames, this );

                _activateLowerToUpper = kernels[0];
                _activateLowerToUpperBinary = kernels[1];
                _activateUpperToLower = kernels[2];
                _accumulateErrors = kernels[3];
                _sumErrorsDetector = kernels[4];
                _sumErrorsRecon = kernels[5];
                _clearDetectorErrors = kernels[6];
                _updateWeights = kernels[7];   
    
                //_accumulateErrors = CreateKernel(program, "AccumulateErrors");
                //_clearDetectorErrors = CreateKernel(program, "ClearDetectorErrors");

                //// cj: Moved constants out of CL code to make it more concise and expressive and easier to touch
                //// Since a recompile was required, I just used comments here with different (albiet poorly named) kernels
                //// At least these can be switch at runtime in the future with litte effort
                //_activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary_Binary" );
                ////_activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary_Linear" );
                ////_activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper_Binary" );
                //_activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper_Linear" );

                ////_activateLowerToUpperBinary = CreateKernel( program, "ActivateLowerToUpperBinary" );
                ////_activateLowerToUpper = CreateKernel( program, "ActivateLowerToUpper" );

                //_activateUpperToLower = CreateKernel( program, "ActivateUpperToLower" );
                //_updateWeights = CreateKernel(program, "UpdateWeights");
                //_sumErrorsDetector = CreateKernel( program, "SumErrorsDetector" );
                //_sumErrorsRecon = CreateKernel( program, "SumErrorsRecon" );       
            }
 
	    #endregion

        #region - Program Functionality -

            public TrainingError Train( ComputationEngine engine )
            {
                float[] errorSums = new float[2];
                TrainingError trainingError = new TrainingError();

                List<IMem> randomizedTestIMemCases = _testCaseMems.ToList();
                randomizedTestIMemCases.Shuffle();

                WriteTrainingRate( engine );

                randomizedTestIMemCases.BatchForEach( BatchSize, testCases => {

                    engine.Flush();

                    ClearDetectorErrors( engine );

                    GetTestCaseErrors( engine, testCases );

                    UpdateWeights( engine );

                    SumDetectorErrors( engine );
                    SumReconstructionErrors( engine );

                    ReadErrorSums( engine, errorSums );

                    trainingError.FeatureDetectorError += errorSums[0];
                    trainingError.ReconstructionError += errorSums[1];
                } );

                ReadWeights( engine );

                // Make sure we're done with everything that's been requested before
                engine.Flush();

                // TODO: Wait for the event that should be returned by _errorSumsIMem.ReadFloatArray...                
                return trainingError;

            }


            private void GetTestCaseErrors( ComputationEngine engine, IEnumerable<IMem> testCases )
            {
                foreach( var testCase in testCases )
                {
                    ActivateLowerToUpperBinary( engine, testCase );

                    ActivateUpperToLower( engine );

                    ActivateLowerToUpper( engine );

                    AccumulateErrors( engine, testCase );
                }
            }

            public void WriteWeights( ComputationEngine engine )
            {
                engine.WriteFloatArray( _weightsMem, _weights.Weights, BlockModes.Blocking );
            }


            private void ReadWeights( ComputationEngine engine )
            {
                // copy results from device back to host
                engine.ReadFloatArray( _weightsMem, _weights.Weights, BlockModes.Blocking );
            }
        
            private void WriteTrainingRate( ComputationEngine engine )
            {
                _epsilon[0] = _weights.GetAdjustedAndScaledTrainingRate( _testCount );
                engine.WriteFloatArray( _epsilonMem, _epsilon, BlockModes.Blocking );
            }

            private void ReadErrorSums( ComputationEngine engine, float[] errorSums )
            {
                engine.ReadFloatArray( _errorSumsMem, errorSums, BlockModes.Blocking );
            }

            private void SumReconstructionErrors( ComputationEngine engine )
            {
                engine.Enqueue( _sumErrorsRecon
                    .SetKernelArgFloat( 0, _detectorErrorMem )
                    .SetKernelArgFloat( 1, _reconstructionErrorMem )
                    .SetKernelArgFloat( 2, _errorSumsMem )
                    .SetKernelArgFloat( 3, _twoPhaseMem ),
                    _weights.LowerLayerSize
                );

            }

            private void SumDetectorErrors( ComputationEngine engine )
            {
                engine.Enqueue( _sumErrorsDetector
                    .SetKernelArgFloat( 0, _detectorErrorMem )
                    .SetKernelArgFloat( 1, _reconstructionErrorMem )
                    .SetKernelArgFloat( 2, _errorSumsMem )
                    .SetKernelArgFloat( 3, _twoPhaseMem ),
                    _weights.LowerLayerSize 
                );

            }

            private void UpdateWeights( ComputationEngine engine )
            {
                engine.Enqueue( _updateWeights
                    .SetKernelArgFloat( 0, _weightsMem )
                    .SetKernelArgFloat( 1, _detectorErrorMem )
                    .SetKernelArgFloat( 2, _epsilonMem ),
                    _weights.LowerLayerSize,
                    _weights.UpperLayerSize
                );
            }

            private void AccumulateErrors( ComputationEngine engine, IMem testCase )
            {
                engine.Enqueue( _accumulateErrors
                    .SetKernelArgFloat( 0, testCase )
                    .SetKernelArgFloat( 1, _detectorErrorMem )
                    .SetKernelArgFloat( 2, _modelMem )
                    .SetKernelArgFloat( 3, _reconstructedMem )
                    .SetKernelArgFloat( 4, _reconstructedModelMem )
                    .SetKernelArgFloat( 5, _reconstructionErrorMem ),
                    _weights.LowerLayerSize,
                    _weights.UpperLayerSize
                );
            }

            private void ActivateLowerToUpper( ComputationEngine engine )
            {
                engine.Enqueue( _activateLowerToUpper
                    .SetKernelArgFloat( 0, _reconstructedMem )
                    .SetKernelArgFloat( 1, _reconstructedModelMem )
                    .SetKernelArgFloat( 2, _weightsMem )
                    .SetKernelArg( 3, Randomized.Default.Next() ),
                    _weights.UpperLayerSize
                );
            }

            private void ActivateUpperToLower( ComputationEngine engine )
            {
                engine.Enqueue( _activateUpperToLower
                    .SetKernelArgFloat( 0, _reconstructedMem )
                    .SetKernelArgFloat( 1, _modelMem )
                    .SetKernelArgFloat( 2, _weightsMem ),
                    _weights.LowerLayerSize
                );
            }

            private void ActivateLowerToUpperBinary( ComputationEngine engine, IMem testCase )
            {
                engine.Enqueue( _activateLowerToUpperBinary
                    .SetKernelArgFloat( 0, testCase )
                    .SetKernelArgFloat( 1, _modelMem )
                    .SetKernelArgFloat( 2, _weightsMem )
                    .SetKernelArg( 3, Randomized.Default.Next() ),
                    _weights.UpperLayerSize
                );
            }

            private void ClearDetectorErrors( ComputationEngine engine )
            {
                engine.Enqueue( _clearDetectorErrors
                    .SetKernelArgFloat( 0, _detectorErrorMem )
                    .SetKernelArgFloat( 1, _reconstructionErrorMem ),
                    _weights.LowerLayerSize,
                    _weights.UpperLayerSize
                );
            }

        #endregion

        #region - IComputationMemory Implementation -
     
            public void AllocateMemory( ComputationEngine engine )
            {
                IntPtr lowerCount = new IntPtr(_weights.LowerLayerSize);
                IntPtr upperCount = new IntPtr(_weights.UpperLayerSize);
                IntPtr testCount = new IntPtr(_testCount);

                _weightsMem = engine.AllocateFloatArray(_weights.Weights, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _detectorErrorMem = engine.AllocateFloatArray(_detectorError, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _epsilonMem = engine.AllocateFloatArray(_epsilon, MemFlags.CopyHostPtr | MemFlags.ReadOnly);
                _modelMem = engine.AllocateFloatArray(_model, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _reconstructedMem = engine.AllocateFloatArray(_reconstructed, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _reconstructedModelMem = engine.AllocateFloatArray(_reconstructedModel, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _reconstructionErrorMem = engine.AllocateFloatArray(_reconstructionError, MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _errorSumsMem = engine.AllocateFloatArray(new float[2], MemFlags.CopyHostPtr | MemFlags.ReadWrite);
                _twoPhaseMem = engine.AllocateFloatArray( _twoPhaseBuffer, MemFlags.CopyHostPtr | MemFlags.ReadWrite );

                _testCaseMems = new List<IMem>();
                for (int testCase = 0; testCase < _testCount; testCase++)
                {
                    IMem testCaseMem = engine.AllocateFloatArray(_inputs[testCase].Values, MemFlags.CopyHostPtr | MemFlags.ReadOnly);
                    engine.WriteFloatArray( testCaseMem, _inputs[testCase].Values, BlockModes.NonBlocking );
                    _testCaseMems.Add(testCaseMem);
                }
            }

            public int BatchSize { get; set; }

        #endregion
    }
}
